from pathlib import Path
import numpy as np
import torch
import yaml
import math
import os
import time
from mynn.datasets import build_dataset, build_dataloader
from mynn.models import build_model
from mynn.losses import L1Loss
from mynn.utils import save_checkpoint, load_checkpint, get_root_logger, flow2img
from PIL import Image

# Read yaml file.
YAML_PATH = 'options/SpyNet4TOF/spynet.yaml'
with open(YAML_PATH, 'r', encoding='utf-8') as f:
    opt = yaml.load(f, Loader=yaml.SafeLoader)

# Set logger.
log_dir = Path('experiments') / opt['exp_name']
os.makedirs(log_dir, exist_ok=True)
log_file = log_dir / f"{opt['test']['log_file']}"
logger = get_root_logger(log_file=log_file)

# Choose CUDA or CPU.
device = "cuda" if opt['test']['cuda'] else "cpu"

# Set GPU list.
gpu_list = ",".join([str(v) for v in opt['test']['gpu_list']])
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
logger.info(f'gpu_list:[{gpu_list}]')

# Get dataset and dataloader.
test_dataset = build_dataset(dataset_opt=opt['test']['dataset'], phase='test')

test_dataloader = build_dataloader(dataset=test_dataset, opt=opt, phase='test')

# Build model.
model = build_model(opt)
model.to(device)

# Load checkpoint.
model, current_iter = load_checkpint(opt=opt, model=model)

logger.info('Testing start.')

model.eval()
total_step = len(test_dataloader)
save_root = Path('./experiments') / opt['exp_name'] / 'results' / opt['test']['dataset']['type']
logger.info(f'Save path:{save_root}')
for step, data in enumerate(test_dataloader):
    # Unpack data.
    ref = data['ref'].to(device)
    supp = data['supp'].to(device)
    flow = data['flow'].to(device)
    key = data['key'][0]

    # Forward propagation.
    flow_hat = model(ref, supp)
    flow_hat = torch.squeeze(flow_hat).permute(1, 2, 0)
    flow_img = flow2img(flow_hat)

    # Save result.
    clip_name, image_name = key.split('/')
    save_dir = save_root / clip_name
    os.makedirs(save_dir, exist_ok=True)
    save_path = save_dir / f'{image_name}.png'

    flow_img = Image.fromarray(flow_img)
    flow_img.save(save_path)
    print(f'{step+1}/{total_step}: {key}')

logger.info('Test complete.')